import torch as pt
from torch import nn
from msi_utils import cumprod_exclusive
from pytorch_warp_msi import warp_imgs_msi
from torch import optim

MSI_N = 10

sigmoid_offset = 5

class MsiModel(nn.Module):
  def __init__(self, dmin, dmax, msi_resolution_w, msi_resolution_h,
    theta_half_range, phi_half_range, rank):
    super().__init__()
    self.radius = pt.linspace(dmin, dmax, MSI_N).view(-1,1,1).cuda(rank)
    msi_rgba = pt.empty([
      MSI_N,
      1,
      msi_resolution_h,
      msi_resolution_w,
      4]).uniform_(-1, 1)
    self.msi_rgba = nn.Parameter(msi_rgba)

    msi_offset = pt.zeros([
      MSI_N,
      1,
      msi_resolution_h,
      msi_resolution_w,
      1])
    self.msi_offset = nn.Parameter(msi_offset)

    self.theta_half_range = theta_half_range
    self.phi_half_range = phi_half_range

  def forward(self, pixel_rays_targ, R, T, use_offset=True):
    msi_rgba = self.msi_rgba.clone()
    msi_rgba[...,3] -= sigmoid_offset
    rgba_sig = pt.sigmoid(msi_rgba)

    offset = pt.tanh(self.msi_offset)
    warped_rgba = warp_imgs_msi(
      rgba_sig, pixel_rays_targ, R, T, self.radius,
      self.theta_half_range, self.phi_half_range,
      offset=offset if use_offset else None)

    weight = cumprod_exclusive(1 - warped_rgba[...,3])
    return pt.sum(warped_rgba[...,:3] * weight[...,None] * warped_rgba[...,3:],
      dim=0, keepdim=True)

  def get_optim(self, lr):
    return optim.Adam([
      {'params': self.msi_rgba, 'lr': lr},
      {'params': self.msi_offset, 'lr': lr * 0.01}])
